package org.example.service.impl;

import org.example.service.OllamaDelegete;
import org.springframework.ai.chat.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatClient;
import org.springframework.ai.ollama.api.OllamaApi;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.List;
import java.util.stream.Collectors;

@Service
public class OllamaImpl implements OllamaDelegete {

    OllamaApi ollamaApi;
    OllamaChatClient chatClient;
    {
        // 实例化ollama
        ollamaApi = new OllamaApi();
        OllamaOptions options = new OllamaOptions();
        options.setModel("llama3");
        chatClient = new OllamaChatClient(ollamaApi).withDefaultOptions(options);
    }

    /**
     * 普通文本调用
     *
     * @param msg
     * @return
     */
    @Override
    public Object getOllame(String msg) {
        Prompt prompt = new Prompt(msg);
        ChatResponse call = chatClient.call(prompt);
        return call.getResult().getOutput().getContent();
    }

    /**
     * 流式调用
     *
     * @param msg
     * @return
     */
    @Override
    public Object getOllameByStream(String msg) {
        Prompt prompt = new Prompt(msg);
        Flux<ChatResponse> stream = chatClient.stream(prompt);
        List<String> result = stream.toStream().map(a -> {
            return a.getResult().getOutput().getContent();
        }).collect(Collectors.toList());
        return result;
    }

}
